#include <rclcpp/rclcpp.hpp>

// MoveIt Servo
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>

#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>
#include <std_msgs/msg/bool.hpp>
#include <tf2_ros/buffer.h>

#include <cmath>
#include <string>
#include <chrono>
#include <memory>

using namespace std::chrono_literals;

namespace imp_node {

// =================== 可读常量（根据需要改） ===================
// 目标恒力（牛顿）
constexpr double kF_target = 5.0;
// 映射系数：u = kappa * (F_target - F_meas)
constexpr double kKappa    = 1.0;

// 阻抗（法向 X 轴）：m_d * xdd + b_d * xdot + k_d * x = u
constexpr double kMd = 1.0;   // 虚拟质量
constexpr double kBd = 40.0;  // 虚拟阻尼
constexpr double kKd = 0.0;   // 虚拟弹簧（接触前建议 0）

// 限幅（法向速度/加速度）
constexpr double kVmax = 0.40;  // m/s
constexpr double kAmax = 1.50;  // m/s^2

// 恒力达标：|Fref - Fe| < tol 且持续 stable_time
constexpr double kForceTol      = 0.5; // N
constexpr double kStableTimeSec = 0.5; // s

// 扫描：进入后在 ZX 平面 45° 方向恒速（x=z=v/√2）
constexpr double kScanSpeed = 0.05; // m/s

// 发送 Twist 的参考坐标系（需与 Servo 配置一致 & TF 存在）
const std::string kTwistFrame = "base_link";

// 控制周期
constexpr std::chrono::milliseconds kLoopPeriod{10}; // 100 Hz
// ===========================================================

static inline double clamp_scalar(double v, double lo, double hi) {
  return v < lo ? lo : (v > hi ? hi : v);
}

// —— 节点与通信对象 —— //
rclcpp::Node::SharedPtr node_;
rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_pub_;
rclcpp::Publisher<std_msgs::msg::Bool>::SharedPtr scan_flag_pub_;
rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr force_sub_;

planning_scene_monitor::PlanningSceneMonitorPtr psm_;
std::unique_ptr<moveit_servo::Servo> servo_;
std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
rclcpp::TimerBase::SharedPtr timer_;

const rclcpp::Logger LOGGER = rclcpp::get_logger("plane_ardc");

// —— 状态量 —— //
enum class Mode : uint8_t { APPROACH = 0, SCAN = 1 };
double force_meas = 0.0; // 从 /sim_force 读取的 Z 轴力（如需改轴，改回调）
double x  = 0.0;         // 法向位移“虚拟状态”
double xd = 0.0;         // 法向速度
double stable_accum = 0.0;
Mode   mode = Mode::APPROACH;

rclcpp::Time t_prev;
bool has_prev_time = false;

// 订阅力（此处取 Z 轴；如果你的力在别的轴，改这里）
void forceCallback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg) {
  force_meas = msg->wrench.force.z;
  RCLCPP_INFO_THROTTLE(
    LOGGER, *node_->get_clock(), 500,
    "[/sim_force] Fz=%.3f N (target=%.2f)", force_meas, kF_target);
}

// 控制循环（阻抗 + 状态机 + Twist 发布）
void controlTimer() {
  // dt
  const rclcpp::Time now = node_->now();
  double dt = has_prev_time ? std::max(1e-4, (now - t_prev).seconds()) : 0.01;
  t_prev = now;
  has_prev_time = true;

  // 力误差 & 虚拟驱动力
  const double eF = kF_target - force_meas;
  const double u  = kKappa * eF;

  // 二阶阻抗（法向 X 轴）
  double xdd = (u - kBd * xd - kKd * x) / std::max(1e-6, kMd);
  xdd = clamp_scalar(xdd, -kAmax, kAmax);         // 限加速度
  double xd_new = clamp_scalar(xd + xdd * dt,     // 积分并限速
                               -kVmax, kVmax);
  x += xd_new * dt;
  xd = xd_new;

  // 达标累计
  if (std::fabs(eF) < kForceTol) stable_accum += dt;
  else                           stable_accum  = 0.0;

  // 进扫描
  if (mode == Mode::APPROACH && stable_accum >= kStableTimeSec) {
    mode = Mode::SCAN;
    std_msgs::msg::Bool on; on.data = true;
    scan_flag_pub_->publish(on);
    RCLCPP_WARN(LOGGER, "====== SCAN_START：力已稳定，开始 ZX 平面 45° 扫描 ======");
  }

  // 组装 Twist
  geometry_msgs::msg::TwistStamped msg;
  msg.header.stamp = now;
  msg.header.frame_id = kTwistFrame;

  double vx = xd;   // 法向速度：linear.x
  double vz = 0.0;  // 切向速度：linear.z

  const double pitch = -0.7;
  const double cx = kScanSpeed * std::cos(pitch);
  const double cz = -kScanSpeed * std::sin(pitch);

  if (mode == Mode::SCAN) {
    const double c = kScanSpeed * M_SQRT1_2; // v/√2
    vx = clamp_scalar(0 + cx, -kVmax, kVmax);
    vz = clamp_scalar(vz + cz, -kVmax, kVmax);
  }

  const double vx_pub = vx;
  const double vz_pub = vz;

  msg.twist.linear.x  = vx;
  msg.twist.linear.y  = 0.0;
  msg.twist.linear.z  = vz;
  msg.twist.angular.x = 0.0;
  msg.twist.angular.y = 0.0;
  msg.twist.angular.z = 0.0;

  twist_pub_->publish(msg);

  RCLCPP_INFO_THROTTLE(
    LOGGER, *node_->get_clock(), 50,
    "[IMP] mode=%s Fe=%.3fN Fref=%.3fN e=%.3f | x=%.4f xd=%.4f | v_pub(x)=%.4f m/s v_pub(z)=%.4f m/s",
    (mode == Mode::APPROACH ? "APPROACH" : "SCAN"),
    force_meas, kF_target, (kF_target - force_meas), x, xd,
    vx_pub, vz_pub);
}

} // namespace imp_node

int main(int argc, char** argv) {
  using namespace imp_node;

  rclcpp::init(argc, argv);
  try {
    // 用自由函数风格，避免构造期 shared_from_this() 导致 bad_weak_ptr
    rclcpp::NodeOptions opts;
    opts.use_intra_process_comms(false);
    node_ = std::make_shared<rclcpp::Node>("plane_ardc", opts);

    // PSM（多数 Servo 能力依赖）
    tf_buffer_ = std::make_shared<tf2_ros::Buffer>(node_->get_clock());
    psm_ = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
      node_, "robot_description", tf_buffer_, "planning_scene_monitor");
    if (!psm_ || !psm_->getPlanningScene()) {
      RCLCPP_ERROR(LOGGER, "Planning scene 未正确配置（robot_description/TF/joint_states？）");
      throw std::runtime_error("planning_scene not ready");
    }
    psm_->startStateMonitor("/joint_states");
    psm_->setPlanningScenePublishingFrequency(25.0);
    psm_->startPublishingPlanningScene(
      planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE,
      "/moveit_servo/publish_planning_scene");
    psm_->startSceneMonitor();
    psm_->providePlanningSceneService();

    // MoveIt Servo（需外部 YAML 参数）
    auto servo_params = moveit_servo::ServoParameters::makeServoParameters(node_);
    servo_ = std::make_unique<moveit_servo::Servo>(node_, servo_params, psm_);
    servo_->start();

    // 通信
    twist_pub_ = node_->create_publisher<geometry_msgs::msg::TwistStamped>("~/delta_twist_cmds", 10);
    scan_flag_pub_ = node_->create_publisher<std_msgs::msg::Bool>("~/scan_active", 1);

    force_sub_ = node_->create_subscription<geometry_msgs::msg::WrenchStamped>(
      "/sim_force", rclcpp::SensorDataQoS(), &forceCallback);

    // 控制循环
    timer_ = node_->create_wall_timer(kLoopPeriod, &controlTimer);

    rclcpp::executors::SingleThreadedExecutor exec;
    exec.add_node(node_);
    exec.spin();

  } catch (const std::exception& e) {
    RCLCPP_ERROR(LOGGER, "启动失败：%s", e.what());
  }
  rclcpp::shutdown();
  return 0;
}
